#include "Ray.h"
#include "Vector3.h"
#include "miro.h"
#include "AARect.h"
#include <math.h>
#include <stdexcept>
#include <iostream>
#include <fstream>
#include <cstdio>
#include <cstdlib>

#include "mt-64.c"


AARect light = AARect(Vector3(-1.75,-1,-.05), Vector3(1.75, -1, .05), YAXIS, Vector3(0,1,0));
AARect mirror = AARect(Vector3(-2,1,-2), Vector3(2, 1, 2), YAXIS, Vector3(0,-1,0));
AARect plane = AARect(Vector3(-1,0,-1), Vector3(1, 0, 1), YAXIS, Vector3(0,1,0));
AARect plane2 = AARect(Vector3(-2, 0, -2), Vector3(2, 0, 2), YAXIS, Vector3(0,1,0));
double photonPower = 100;
double radius = .25;
double radius2 = radius*radius;
float plarge = .1;
float mlsStep = 1./50.;


float getRand()
{
	//return ((float) rand()) / RAND_MAX;
	return (float)genrand64_real1();
}

struct RandVars
{
	float u1, u2, u3, u4, u5;
	void generate()
	{
		u1 = getRand();
		u2 = getRand();
		u3 = getRand();
		u4 = getRand();
		u5 = getRand();
	}
};

float mutateVar(float u, float mutationSize)
{
	float du1, du2, u1, sgn;
	do {
		du1 = getRand();
		du2 = getRand();
		if (du1 < 0.5)
			sgn = -1;
		else
			sgn = 1;
		u1 = u + sgn * pow(du2, 1 + 1/mutationSize);
	} while (u1 < 0 || 1 < u1);
	//printf("%.3f \t %.3f \t %.8f\n", u, u1, mutationSize);
	return u1;
}

struct RandPath
{
	RandVars u;
	Ray ray;
	Vector3 n;

	RandPath(Vector3 normal=Vector3(0,1,0))
	{
		n = normal;
	}

	void generate()
	{
		u.generate();
		recompute();
	}
	void recompute()
	{
		// origin - this bit is scene specific
		float ox = 3.5 * u.u4 - 1.75;
		float oz = .1 * u.u5 - .05;

		// direction
		float sin2 = u.u1;

		float sintheta = sqrt(sin2);
		float costheta = cos(asin(sintheta));
		float phi = 2 * PI * u.u2;

		float x = -sintheta * sin(phi);
		float y = sintheta * cos(phi);

		float cos2 = 1.0f - (sintheta * sintheta);
		float z = sqrt(cos2);

		Vector3 tangentRay;
		if (fabs(n.x) > fabs(n.y))
			tangentRay = Vector3(0, 1, 0);
		else
			tangentRay = Vector3(1, 0, 0);
		tangentRay = cross(n, tangentRay);
		tangentRay.normalize();

		Vector3 rayDir = x * tangentRay + y * cross(n, tangentRay) + z * n;

		ray = Ray(Vector3(ox, -1, oz), rayDir);
	}
	void mutate(float mutationSize)
	{
		// TODO: idea: maybe only mutate one of these (either cycle or pick randomly)
		// mutate rand vars
		u.u1 = mutateVar(u.u1, mutationSize);
		u.u2 = mutateVar(u.u2, mutationSize);
		u.u3 = mutateVar(u.u3, mutationSize);
		u.u4 = mutateVar(u.u4, mutationSize);
		u.u5 = mutateVar(u.u5, mutationSize);

		recompute();
	}
};

#define NUM_MEAS_POINTS 100
#define IDEAL_R (0.05)
struct PhotonMap
{
	//const int NUM_MEAS_POINTS = 100;
	float initRadius; //[NUM_MEAS_POINTS]

	Vector3 pos[NUM_MEAS_POINTS]; //

	unsigned int newPhotons[NUM_MEAS_POINTS];
	unsigned long long hitPhotons[NUM_MEAS_POINTS];
	long double fracHitPhotons[NUM_MEAS_POINTS];
	unsigned long long totalPhotons;

	double alphas[NUM_MEAS_POINTS];
	double scale[NUM_MEAS_POINTS];
	float radius2[NUM_MEAS_POINTS];
	double power[NUM_MEAS_POINTS];
	float E[NUM_MEAS_POINTS];

	PhotonMap(float initR=1)
	{
		initRadius = initR;
		totalPhotons = 0;
		float step = 2.0 / (NUM_MEAS_POINTS-1); // right now, only do half
		for (int i = 0; i < NUM_MEAS_POINTS; ++i)
		{
			pos[i] = Vector3(-1 + i*step, 0, 0);
			//printf("\t %.3f",pos[i].x);
			newPhotons[i] = 0;
			hitPhotons[i] = 0;
			scale[i] = 0;
			radius2[i] = initRadius*initRadius;
			power[i] = 0;
		}
	}
	~PhotonMap()
	{
	}

	bool isVisible(const Ray &r, HitInfo &hit)
	{
		HitInfo hitPlane, hitMirror;
		// doesn't hit backside of plane, AND hits mirror
		if (!plane.intersect(hitPlane, r) && mirror.intersect(hitMirror, r))
		{
			Vector3 reflection = r.d - 2 * (dot(r.d, hitMirror.N)) * hitMirror.N;
			reflection.normalize();
			Ray r2 = Ray(hitMirror.P, reflection);
			// hits top side of plane
			if(plane2.intersect(hit, r2))
			{
				// if ray hit any measurement point's disc
				for (int i = 0; i < NUM_MEAS_POINTS; ++i)
				{
					if ((hit.P - pos[i]).length2() < radius2[i])
					{
						return true;
					}
				}
			}
		}
		return false;
	}

	bool isVisible(const RandPath &p, HitInfo &hit)
	{
		return isVisible(p.ray, hit);
	}

	bool isVisible(Ray &r)
	{
		HitInfo hit;
		return isVisible(r, hit);
	}

	bool isVisible(RandPath &p)
	{
		return isVisible(p.ray);
	}

	void addPhoton(const Vector3 &photonPos)
	{
		for (int i = 0; i < NUM_MEAS_POINTS; ++i)
		{
			if ((photonPos - pos[i]).length2() < radius2[i])
			{
				newPhotons[i]++;
			}
		}
		totalPhotons++;
	}
	void addPhoton(const Ray &r)
	{
		HitInfo hit;
		if (isVisible(r, hit))
		{
			for (int i = 0; i < NUM_MEAS_POINTS; ++i)
			{
				if ((hit.P - pos[i]).length2() < radius2[i])
				{
					newPhotons[i]++;
				}
			}
		}
		// increase total count either way
		// it's responsibility of appm to make sure this only gets called when isVisible
		totalPhotons++;
	}
	void addPhoton(const RandPath &p)
	{
		addPhoton(p.ray);
	}

	void splat(float alpha, float photonPower)
	{
		// update all measurement points
		for (int i = 0; i < NUM_MEAS_POINTS; ++i)
		{
			// TODO: could try adapting alpha
			//alpha = alphas[i] = 1 - 2*(radius2[i]-IDEAL_R) / (radius2[i]-IDEAL_R + 1);
			int m = newPhotons[i];
			if (m > 0)
			{
				//scale[i] = (alpha * hitPhotons[i] + alpha * m) / (alpha * hitPhotons[i] + m);
				scale[i] = (fracHitPhotons[i] + alpha * m) / (fracHitPhotons[i] + m);
				radius2[i] *= scale[i];
				//hitPhotons[i] += m;
				fracHitPhotons[i] += alpha * m;
				power[i] = (power[i] + m * photonPower) * scale[i]; // scaled by PI - cancelled out BRDF
			}
			// reset new photon counter
			newPhotons[i] = 0;
		}
	}

	void printE(float normFactor)
	{
//		for (int i = 0; i < NUM_MEAS_POINTS; ++i)
//		{
//			E[i] = normFactor * power[i] / totalPhotons / (PI * radius2[i]);
//			printf("\t %.3f", E[i]);
//		}

		float E0 = normFactor * power[0] / totalPhotons / (PI * radius2[0]);
		float E49 = normFactor * power[49] / totalPhotons / (PI * radius2[49]);
		float E99 = normFactor * power[99] / totalPhotons / (PI * radius2[99]);
		printf("%lld \t %.3f:%.3f:%.3f \t %.3f:%.3f:%.3f \t %.3f:%.3f:%.3f\n", totalPhotons, E0, radius2[0], pos[0].x, E49, radius2[49], pos[49].x, E99, radius2[99], pos[99].x);
	}

	float getE(int i, float normFactor)
	{
		E[i] = normFactor * power[i] / totalPhotons / (PI * radius2[i]);
		return E[i];
	}

};

// cosine-weighted random ray
static Vector3 randDir(Vector3 n, float &sintheta, float &costheta, float u1=-10., float u2=-10.)
{
	if(u1<-9.)
		u1 = getRand();
	if(u2<-9.)
		u2 = getRand();

	float sin2 = u1;

	sintheta = sqrt(sin2);
	costheta = cos(asin(sintheta));
	float phi = 2 * PI * u2;

	float x = -sintheta * sin(phi);
	float y = sintheta * cos(phi);

	float cos2 = 1.0f - (sintheta * sintheta);
	float z = sqrt(cos2);

	Vector3 tangentRay;
	if (fabs(n.x) > fabs(n.y))
		tangentRay = Vector3(0, 1, 0);
	else
		tangentRay = Vector3(1, 0, 0);
	tangentRay = cross(n, tangentRay);
	tangentRay.normalize();

	return x * tangentRay + y * cross(n, tangentRay) + z * n;
}

// uniform random ray
static Vector3 randUDir(Vector3 n, float &sintheta, float &costheta)
{
	float x, y, z;

	do
	{
		x = ((float) rand()) / RAND_MAX;
		y = ((float) rand()) / RAND_MAX;
		z = ((float) rand()) / RAND_MAX;
	} while (x * x + y * y + z * z > 1);

	Vector3 dir(x, y, z);
	dir.normalize();

	Vector3 tangentRay;
	if (fabs(n.x) > fabs(n.y))
		tangentRay = Vector3(0, 1, 0);
	else
		tangentRay = Vector3(1, 0, 0);
	tangentRay = cross(n, tangentRay);
	tangentRay.normalize();

	Vector3 result = dir.x * tangentRay + dir.y * cross(n, tangentRay) + dir.z
			* n;

	costheta = dot(n, dir);
	sintheta = sin(acos(costheta));

	return result;
}

bool pathTraceRay(Ray &r, HitInfo &hit) {
	if (mirror.intersect(hit, r))
	{
		Vector3 reflection = r.d - 2 * (dot(r.d, hit.N)) * hit.N;
		reflection.normalize();
		Ray r2 = Ray(hit.P, reflection);
		if (!plane.intersect(hit, r2) && light.intersect(hit, r2))
		{
			return true;
		}
	}
	return false;
}

// Task 1: Path Tracing
float pathTraceScene(int totalRays, bool verbose = false, bool silent = false)
{
	int numRays = totalRays/100;
	int printInterval = 10000;
	int outputInterval = 1000;
	int countHits = 0;
	float EVals[100] = {0.0};
//	AARect light = AARect(Vector3(-1.75,-1,-.05), Vector3(1.75, -1, .05), YAXIS, Vector3(0,1,0));
//	AARect mirror = AARect(Vector3(-2,1,-2), Vector3(2, 1, 2), YAXIS, Vector3(0,-1,0));
//	AARect plane = AARect(Vector3(-1,0,-1), Vector3(1, 0, 1), YAXIS, Vector3(0,1,0));
	
	float P = 100; // light power
	float A = 3.5 * .1; // light area
	float L = P / A / PI; // Radiance per ray - figure this out later
	double E;
	double totalE = 0;
	double totalL = 0;
	float sintheta, costheta;

	std::ofstream f("pt.txt");

	for(int j = 0; j<100; j++) {
		totalE = 0;
		totalL = 0;

		float point = j/49.5 - 1.;
		Vector3 rayO = Vector3(point,0,0);
		for (int i = 1; i <= numRays; i++)
		{
			Vector3 randRay = randDir(Vector3(0, 1, 0), sintheta, costheta);
			Ray r = Ray(rayO, randRay);
			HitInfo hit;
			if(pathTraceRay(r, hit))
			{
				countHits++;
				totalE += L * PI;
				totalL += L;
			}


			if(!verbose && i % (numRays/10) == 0)
			{
				//fprintf(stderr,"Completed: %.1f%%\r", (float) j + ((double) i)/((double) numRays));
			}
			else if (verbose && i % printInterval == 0)
			{
				E = totalE / (i + 1);
				printf("%d\t%.8f\n", i, E);
			}

			if(i % outputInterval == 0)
			{
				E = totalE / (i + 1);
				f << E << ", ";
			}

			EVals[j] = totalE / (i + 1);
		}
		if(!verbose && !silent) {
			fprintf(stderr, "Completed: %d%%\r", j+1);
		}
		f << E << "\n";
	}

	f.close();

	E = totalE / numRays;
	if (verbose)
	{
		printf("%d\t%.5f\n", numRays, E);
		printf("L = %.5f\n", totalL / numRays);
	}
	else
	{
		//printf("Completed: 100.0%%\n");
	}

	float b = 0.;
	for(int i=0; i<100; i++) {
		b += EVals[i]/100.;
	}

	return b;
}

// Task 2: Progressive Photon Mapping
float ppmScene(bool verbose = false)
{

	float alpha = 0.7;
	float radius = 1.0;
	unsigned int photonsPerSplat = 1;
	unsigned long long maxPhotons = 100000000;

	//photon map!
	PhotonMap photonMap(radius);

	std::ofstream f("ppm.txt");

	int printInterval = 1000000;
	int outputInterval = 1000000;

	float E, V; // nomalization factor
	RandPath randPath;

	while (photonMap.totalPhotons < maxPhotons)
	{
		for (int j = 0; j < photonsPerSplat; j++)
		{
			randPath.generate();
			photonMap.addPhoton(randPath);
		}
		photonMap.splat(alpha, photonPower);


		if(photonMap.totalPhotons % outputInterval == 0)
		{
			bool printToFile = false;
			if (photonMap.totalPhotons == 1000000 || photonMap.totalPhotons == 10000000 || photonMap.totalPhotons == 100000000)
				printToFile = true;

			if (printToFile)
			{
				for (int i=0; i<NUM_MEAS_POINTS; ++i)
				{
					E = photonMap.getE(i, 1);
						f << "\t" << E;
				}
				f << "\n";
			}
			photonMap.printE(1);
		}
	}

	printf("Completed: 100.0%%\n");
	f << "\n";
	f << alpha << "\t" << photonMap.initRadius << "\t" << photonsPerSplat << "\n";
	f.close();

	return E;
}


// this creates a new RandPath from another one, which does not get affected
// (we don't want to mess up the current path)
RandPath mutatePath(const RandPath &p, float mutationSize)
{
	RandPath newPath = p;
	newPath.mutate(mutationSize);
	return newPath;
}

float appmScene(bool verbose = false)
{
	float alpha = 0.9;
	float radius = 0.25;
	unsigned int photonsPerSplat = 1000000;
	unsigned long long maxPhotons = 100000000;

	//photon map!
	PhotonMap photonMap(radius);

	std::ofstream f("appm.txt");

	int printInterval = 1000000;
	int outputInterval = 1000000;

	float E, V; // nomalization factor

	// appm params
	float mutationSize = 1;
	unsigned long long acceptedCount = 1;
	unsigned long long mutatedCount = 0;
	unsigned long long uniformCount = 1;
	float acceptRatio = 1;
	bool accepted = false;
	RandPath currentPath, uniformPath, candidatePath;
	do
	{
		currentPath.generate();
	} while (!photonMap.isVisible(currentPath));
	while (photonMap.totalPhotons < maxPhotons)
	{
		for (int i = 0; i < photonsPerSplat; i++)
		{
			uniformPath.generate();
			if (photonMap.isVisible(uniformPath))
			{
				// we found a new good path - set it to current
				currentPath = uniformPath;
				uniformCount++;
			}
			else
			{
				// mutate current path
				candidatePath = mutatePath(currentPath, mutationSize);
				mutatedCount++;
				if (photonMap.isVisible(candidatePath))
				{
					currentPath = candidatePath;
					acceptedCount++;
				}
				acceptRatio = float(acceptedCount) / mutatedCount;
				mutationSize = mutationSize + (acceptRatio - 0.234)/mutatedCount;
			}
			// accumulate photon from currentPath
			// we'll splat later
			photonMap.addPhoton(currentPath);
		}

		// splat
		photonMap.splat(alpha, photonPower);

		// display
		if(photonMap.totalPhotons % outputInterval == 0)
		{
			bool printToFile = false;
			if (photonMap.totalPhotons == 1000000 || photonMap.totalPhotons == 10000000 || photonMap.totalPhotons == 100000000)
				printToFile = true;

			V = float(uniformCount) / photonMap.totalPhotons;
			if (printToFile)
			{
				for (int i=0; i<NUM_MEAS_POINTS; ++i)
				{
					E = photonMap.getE(i, V);
						f << "\t" << E;
				}
				f << "\n";
			}
			photonMap.printE(V);
		}
	}

	printf("Completed: 100.0%%\n");
	f << "\n";
	f << alpha << "\t" << photonMap.initRadius << "\t" << photonsPerSplat << "\n";
	f.close();

	return E;
}

float mutateValue( float value, float minv = 0.0, float maxv = 1.0 )
{
	float s1 = 1./1024, s2 = 1./64;
	float dv = s2*exp(-log(s2/s1)*getRand());

	if (getRand() < 0.5)
	{
		value += dv;
		if (value > maxv)
			value -= (maxv - minv);
	}
	else
	{
		value -= dv;
		if (value < minv)
			value += (maxv - minv);
	}

	return value;
}

struct mlsPath
{
	float u1, u2;
	int ind;
	Ray ray;
	Vector3 n;

	mlsPath(Vector3 normal=Vector3(0,1,0))
	{
		n = normal;
	}

	void generate()
	{
		float sintheta, costheta;
		u1 = getRand();
		u2 = getRand();
		//ind = index;

		//ind = (int) 100*genrand64_real2();
		//float point = ind/49.5 - 1.;
		float point = 2.*getRand()-1.;
		ind = std::min((int) ((point + 1.)*50.), 99);

		Vector3 randRay = randDir(Vector3(0, 1, 0), sintheta, costheta, u1, u2);
		Vector3 rayO = Vector3(point, 0,0);
		ray = Ray(rayO, randRay);
	}
	void recompute()
	{
		float sintheta, costheta;
		Vector3 randRay = randDir(n, sintheta, costheta, u1, u2);

		ray = Ray(ray.o, randRay);
	}
	void mutate()
	{
		u1 = mutateValue(u1);
		u2 = mutateValue(u2);

		ray.o.x = mutateValue( ray.o.x, -1., 1. );
		ind = std::min((int) ((ray.o.x + 1.)*50.), 99);
	}
};

float mlsScene(bool verbose = false)
{
	int M = 100000000;

	float P = 100;
	float A = 3.5 * .1;
	float L = P / A / PI;
	double E =0;
	double EVals[100] = {0.};
	double allVals[3][100] = {0.};

	int counts[100] = {0};
	mlsPath curPath = mlsPath();
	mlsPath newPath = mlsPath();
	float a;
	plarge = 0.;

	float b = pathTraceScene(20000000, false, true);

	int hits =0, misses=0;
	bool foundPath = false;
	while (!foundPath)
	{
		curPath.generate();
		Ray r = curPath.ray;
		HitInfo hit;

		if (pathTraceRay(r, hit))
			foundPath = true;
	}

	int countHits = 0;
	for(int i=0; i<M; i++)
	{
		if(getRand() <= .6)
		{
			newPath.generate();
			plarge = 0.;
		}
		else
		{
			plarge += .01;
			newPath.u1 = curPath.u1;
			newPath.u2 = curPath.u2;
			newPath.n = curPath.n;
			newPath.ray = curPath.ray;
			newPath.ind = curPath.ind;

			newPath.mutate();
			newPath.recompute();
		}
		HitInfo hit;
		if(pathTraceRay(newPath.ray, hit))
			a = 1.;
		else
			a = 0.;

		counts[newPath.ind]++;

		EVals[newPath.ind] += 100.*a*b/((double) M);
		EVals[curPath.ind] += 100.*(1.-a)*b/((double) M);

		if(getRand() <= a ) {
			curPath = newPath;
		}

		if(i%(M/100) == 0)
			fprintf(stderr, "%d%%\r", i/(M/100)+1 );

		if(i+1==10000000) {
			for(int k=0; k<100; k++) {
				allVals[0][k] = EVals[k];
			}
		}
		if(i+1==100000000) {
			for(int k=0; k<100; k++)
				allVals[1][k] = EVals[k];
		}
		if(i+1==1000000000) {
			for(int k=0; k<100; k++)
				allVals[2][k] = EVals[k];
		}

	}

	std::ofstream f("mls.txt");
	for(int q=0; q<3; q++) {
		for(int k=0; k<100; k++)
			f << allVals[q][k] << "\t";
		f << "\n";
	}

	f.close();
	for(int i=0; i<10; i++)
	{
		for(int j=0; j<10; j++)
		{
			printf("%.3f, ", EVals[i*10 +j]);
		}
		printf("\n");
	}


	return E;
}

int main(int argc, char*argv[])
{
	float E;
  	
	//E = pathTraceScene(2000000, false);
	//E = ppmScene(false);
	E = appmScene(false);
	//E = mlsScene(true);

	printf("E = %.8f\n", E);
	return 0;
}

